package com.dayang.framework.biz.limiter.aspect;

import cn.hutool.core.util.RandomUtil;
import com.dayang.framework.common.exception.BizException;
import com.dayang.framework.common.exception.CommonResponseCodeEnum;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;

import java.util.Map;

/**
 * @author: dayang
 * @since: 2024/6/14
 * @description:
 */
@Aspect
@Slf4j
public class MyRateLimiterAspect {
    /**
     * 以资源名称为key值，存储多个的限流器
     */
    private Map<String, RateLimiter> limiterMap = Maps.newConcurrentMap();

    @Around("@annotation(myRateLimiter)")
    public Object process(ProceedingJoinPoint point,MyRateLimiter myRateLimiter) throws Throwable {
        String key = RandomUtil.randomString(6);
        RateLimiter rateLimiter = this.getRateLimiter(key,myRateLimiter);
        boolean acquired = rateLimiter.tryAcquire(
                myRateLimiter.permits(),
                myRateLimiter.timeout(),
                myRateLimiter.timeUnit()
        );
        if (!acquired){
            throw new BizException(CommonResponseCodeEnum.VERIFICATION_CODE_SEND_FREQUENTLY);
        }
        return point.proceed();
    }

    /**
     * 获取 RateLimiter对象，如果容器中没有对应的资源的限流器，先创建一个限流器
     * @param key
     * @param myRateLimiter
     * @return
     */
    private RateLimiter getRateLimiter(String key, MyRateLimiter myRateLimiter) {
        RateLimiter rateLimiter = limiterMap.get(key);
        if (rateLimiter == null){
            rateLimiter = this.createRateLimiter(myRateLimiter);
        }
        return rateLimiter;
    }

    /**
     * 新建限流器，根据参数分别创建平滑预热限流器或者是平滑突发限流器
     * @param myRateLimiter
     * @return
     */
    private RateLimiter createRateLimiter(MyRateLimiter myRateLimiter) {
        RateLimiter rateLimiter;
        if (myRateLimiter.isWarmup()) {
            rateLimiter = RateLimiter.create(
                    myRateLimiter.permitsPerSecond(),
                    myRateLimiter.warmupPeriod(),
                    myRateLimiter.warmupTimeUnit()

            );
        }else {
            rateLimiter = RateLimiter.create(myRateLimiter.permitsPerSecond());
        }
        return rateLimiter;
    }

}
